#include <iostream>
#include <ostream>
#include <vector>
#include <fstream>
#include <string>
#include <sstream>

#include "app.h"


#include "CImg.h"
using namespace cimg_library;

/* ---------------------------------------------------------------------------
 * Display Variables
 * -------------------------------------------------------------------------*/
CImgDisplay         disp_main;
CImgDisplay         disp_help;
CImg<unsigned char> canvas_main;
CImg<unsigned char> canvas_help;


/* ---------------------------------------------------------------------------
 * Chart junk variables and parameters
 * --------------------------------------------------------------------------*/
std::vector<std::pair<double,double>>* linechart = NULL; 
std::vector<std::pair<double,double>> linechart_in;

CImg<unsigned char> img_in;
CImg<unsigned char> img_rgb, img_grd;
CImg<unsigned char>* img_bg=NULL;
int Dx, Dy;


/* ---------------------------------------------------------------------------
 * Watershed variables and parameters
 * --------------------------------------------------------------------------*/
int blur = 2;
CImg<int>           seg_label;
CImg<unsigned char> seg_face, seg_boundary, grd_norm, marker;
std::vector<int>    selectedId;


/* ----------------------------------------------------------------------------
 * Particl Swarm Optimization variables and parameters
 * ----------------------------------------------------------------------------*/
int Dim = 4; // tx,ty,sx,sy (number of variables)
int numParticles;
int numIterations;
std::vector<double> BoundPosMin(Dim);
std::vector<double> BoundPosMax(Dim);
std::vector<double> BoundVeloMin(Dim);
std::vector<double> BoundVeloMax(Dim);


/* ----------------------------------------------------------------------------
 * User preference and PSO constrains
 * ----------------------------------------------------------------------------*/
int numConstrains = 2;  // unused

// region constrain
double region_factor = 0.5;
CImg<int>* region = NULL;

// scale penalty
double user_scalex = 0.5;  // 0.7% of width of image
double user_scaley = 0.5;  // 0.5% of height of image
double user_tolx   = 2.0;  // tolerance of user preferred scale
double user_toly   = 2.0;  // tolerance of user preferred scale

// edge constrain
double edge_factor = 0.0;
double edge_threshold = 100;

// color constrain
double color_factor = 0.0;
unsigned int nColor = 4;
std::vector<int> clr_picked;
std::vector<int> loc_picked;
std::vector<std::vector<double>> gkernel;
int kernel_radius = 5;

//----------------------------------------------------------------------------
// PSO parameters
//----------------------------------------------------------------------------
// - All parameter values are set to values recommended b a research paper
//   that investigated the effects of various PSO parameter values on a set of
//   benchmark minimization problems.
// - Alternatively, the inertia weight can be dynamically vary according to 
//   the iterations.
// - If you set the cognitive value to be larger than the value of social weight
//   you place more weight on a particle's best known position that on the 
//   swarm's lobal best known position, and vice versa.
//----------------------------------------------------------------------------
double  w = 0.729;    // inertia weight
double c1 = 1.49445;  // cognitive weight
double c2 = 1.49445;  // social weight
double c3 = 0.5*1.49445;  // global weight
// ----------------------------------------------------------------------------




void update_numParticles(int n) {
	numParticles = (n<10) ? 10: n;
}
void update_numIterations(int n) {
	numIterations = (n<1) ? 1000: n; 
}
void update_solution_space(
	int type,
	double bound_pos_min, double bound_pos_max,
	double bound_velo_min, double bound_velo_max
	)
{
	BoundPosMin[type] = bound_pos_min;
	BoundPosMax[type] = bound_pos_max;
	BoundVeloMin[type] = bound_velo_min;
	BoundVeloMax[type] = bound_velo_max;
}


void init_watershed(){
	seg_label   .assign(Dx,Dy,1,1,-1);
	seg_face    .assign(Dx,Dy,1,3, 0);
	seg_boundary.assign(Dx,Dy,1,1, 0);

	CImg<unsigned char> img_gray = img_in.get_resize(-100,-100,-100,1);
	CImg<unsigned char> img_blur = img_gray.get_blur(blur);
	
	CImgList<> grdList = img_blur.get_gradient("xy",2);
	CImg<unsigned char> grd = (grdList[0].pow(2) + grdList[1].pow(2)).sqrt().normalize(0,255);
	grd_norm = grd.get_normalize(10,255);

	//grdList = img_rgb.get_gradient("xy",2);
	//img_grd = (grdList[0].pow(2) + grdList[1].pow(2)).sqrt().normalize(0,255);
	img_grd = grd_norm.get_resize(-100,-100,-100,3);

	marker.assign(Dx,Dy,1,1,0);
}

void update_img_grd() {
	cimg_forXYC(img_grd, x,y,c) {
		if ( (double)grd_norm(x,y) > edge_threshold )
			img_grd(x,y,c) = (double)grd_norm(x,y);
		else
			img_grd(x,y,c) = 0.0;
	}
}


void set_normalmode()
{
	disp_main.set_title("Line junk-chart making: normal mode");
	//canvas_main.assign(img_rgb);
	canvas_main.assign(*img_bg);
}
void set_segmentationmode()
{
	disp_main.set_title("Line junk-chart making: segmentation mode");
	//canvas_main.assign(img_rgb);
	canvas_main.assign(*img_bg);
	cimg_forXYC(canvas_main,x,y,c) {
		//if (seg_face(x,y,c) == 0)
		//	continue;
		canvas_main(x,y,c) = canvas_main(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
	}
	cimg_forXY(canvas_main,x,y) {
		if (seg_boundary(x,y) == 255) {
			canvas_main(x,y,0) = 255; 
			canvas_main(x,y,1) = 255; 
			canvas_main(x,y,2) = 0;
		}
	}
}
void set_selectionmode() {
	disp_main.set_title("Line junk-chart making: selection mode");
	//canvas_main.assign(img_rgb);
	canvas_main.assign(*img_bg);
	cimg_forXYC(canvas_main,x,y,c) {
		bool bSelected = false;
		for (unsigned int i=0; i<selectedId.size(); i++) {
			if (selectedId[i] == seg_label(x,y)) {
				bSelected = true;
				break;
			}
		}

		if (bSelected) continue;

		canvas_main(x,y,c) = canvas_main(x,y,c)*0.7f + seg_face(x,y,c)*0.3f;
	}
	cimg_forXY(canvas_main,x,y) {
		if (seg_boundary(x,y) == 255) {
			canvas_main(x,y,0) = 255; 
			canvas_main(x,y,1) = 255; 
			canvas_main(x,y,2) = 0;
		}
	}
}
void change_viewmode(int mode) {
	switch (mode)
	{
	case 0: set_normalmode(); break;
	case 1: set_segmentationmode(); break;
	case 2: set_selectionmode(); break;
	default: break;
	}
}


void dlmread_linechart(
	const char* fn,
	std::vector<std::pair<double,double>>& l)
{
	std::fstream file(fn, std::ios_base::in);

	int ndata;
	file >> ndata;

	l.resize(ndata);
	for (int i=0; i<ndata; i++) {
		file >> l[i].first;
		file >> l[i].second;
	}

	file.close();
}
void adjust_linechart(
	int width, int height,
	std::vector<std::pair<double,double>>& l
	)
{
	int ndata = (int)l.size();
	double 
		mnx = l[0].first,
		mny = l[0].second,
		mxx = l[ndata-1].first,
		mxy = l[ndata-1].second;

	for (int i=0; i<ndata; i++) {
		double y = l[i].second;
		if (mny > y) mny = y;
		if (mxy < y) mxy = y;
	}

	for (int i=0; i<ndata; i++) {
		double
			x = (l[i].first - mnx)*((double)width/(mxx - mnx)),
			y = (l[i].second - mny)*((double)height/(mxy - mny));

		y = (double)height - y - 1.0;

		l[i].first = x;
		l[i].second = y;
	}

	double
		tx = l[0].first, ty = l[0].second;

	for (int i=0; i<ndata; i++) {
		l[i].first = l[i].first - tx;
		l[i].second = l[i].second - ty;
	}
}


void init_gkernel()
{
	int R = kernel_radius;

	gkernel.resize(2*R+1);
	for (unsigned int i=0; i<gkernel.size(); i++)
		gkernel[i].resize(2*R+1);

	// Create 2d-gaussian mask
	double sigma = 0.3*((2*R+1)/2 - 1) + 0.8;
	double sum = 0.0;
	for (unsigned int y=0; y<gkernel.size(); y++) {
		for (unsigned int x=0; x<gkernel[y].size(); x++)
		{
			double fx = (double)x - (double)R;
			double fy = (double)y - (double)R;
			double w = std::exp(-(fx*fx + fy*fy)/(2.0*sigma*sigma));

			sum += w;
			gkernel[y][x] = w;
		}
	}

	for (unsigned int y=0; y<gkernel.size(); y++) {
		for (unsigned int x=0; x<gkernel[y].size(); x++) {
			gkernel[y][x] /= sum;
		}
	}
}

void refresh_canvas_help()
{
	const unsigned char green[] = {0,255,0}, black[] = {0,0,0};

	std::stringstream ss;
	ss 
	<< "-----------------------------------------------\n"
	<< "Line junk-chart making\n"
	<< "                            by Sungsoo Ha\n"
	<< "-----------------------------------------------\n"
	<< "F1 : Switch display mode\n"
	<< "F2 : Watershed segmentation\n"
	<< "F3 : Select segments\n"
	<< "F4 : Run PSO\n"
	<< "F5 : Switch background\n"
	<< "F6 : Pick color\n"
	<< "-----------------------------------------------\n"
	<< "Scale x [q/w]:   " << user_scalex << "\n" 
	<< "Scale y [a/s]:   " << user_scaley << "\n"
	<< "Scale Tolerence x [e/r]:   " << user_tolx << "\n"
	<< "Scale Tolerence y [d/f]:   " << user_toly << "\n"
	<< "-----------------------------------------------\n"
	<< "Region factor [1/2]:   " << region_factor << "\n"
	<< "-----------------------------------------------\n"
	<< "Edge factor [3/4]:   " << edge_factor << "\n"
	<< "Edge threshold [z/x]:   " << edge_threshold << "\n"
	<< "-----------------------------------------------\n"
	<< "color factor [5/6]:   " << color_factor << "\n"
	<< "color picked:\n";
	canvas_help.assign(200,400,1,3,0).draw_text(5,5,ss.str().c_str(),green);
	//canvas_help.draw_rectangle(5,315,15,325,black)
	//		   .draw_rectangle(20,315,30,325,black)
	//		   .draw_rectangle(35,315,45,325,black)
	//		   .draw_rectangle(50,315,60,325,black);

	unsigned int j = 0;
	for (unsigned int i = 0; i < clr_picked.size()/3; i++) {
		unsigned char clr[3];
		clr[0] = clr_picked[3*i + 0];
		clr[1] = clr_picked[3*i + 1];
		clr[2] = clr_picked[3*i + 2];
		canvas_help.draw_rectangle(5 + i*15, 315 + j*15, 15 + i*15, 325 + j*15, clr);

		j = i/14;
	}

	canvas_help.display(disp_help);
}


int main(int argc, char ** argv)
{
	int mode = 0; // 0: normal, 1: segmentation, 2:selection
	int onhelp = 1; // 0: off, 1: on 
	int bgmode = 0; // 0: rgb, 1: gradient

	char fn[] = ".\\bolt\\bolt_hflip.bmp";
	char line_fn[] = ".\\bolt\\100m.txt";
	
	img_in.assign(fn);
	img_rgb = img_in.get_resize(-100,-100,-100,3);
	img_bg = &img_rgb;
	Dx = img_in.width();
	Dy = img_in.height();

	dlmread_linechart(line_fn, linechart_in);
	adjust_linechart(Dx,Dy,linechart_in);

	// -----------------------------------------------------------------------
	// Default setting for PSO
	// -----------------------------------------------------------------------
	update_numParticles(50);
	update_numIterations(100);
	update_solution_space(0, 0, img_in.width(), -1.0 * img_in.width(), 1.0 * img_in.width());
	update_solution_space(1, 0, img_in.height(), -1.0 * img_in.height(), 1.0 * img_in.height());
	update_solution_space(2, 0.1, 1.0, -1.0, 1.0);
	update_solution_space(3, 0.1, 1.0, -1.0, 1.0);

	linechart = new std::vector<std::pair<double,double>>(linechart_in.size());

	//numConstrains = 2;
	//region_factor = 1.0;
	region = new CImg<int>(Dx,Dy,1,1);

	// -----------------------------------------------------------------------
	// Default setting for Watershed
	// -----------------------------------------------------------------------
	init_watershed();

	disp_main = CImgDisplay(cimg_fitscreen(Dx,Dy,1), "Line junk-chart making: normal mode");	
	canvas_main.assign(*img_bg);
	canvas_main.display(disp_main);

	init_gkernel();
	update_img_grd();

	refresh_canvas_help();

	// -----------------------------------------------------------------------
	// Run Main Program
	// -----------------------------------------------------------------------
	while( !disp_main.is_closed() ) {
		disp_main.wait();

		if (disp_main.key()) {
			switch (disp_main.key())
			{
			case cimg::keyF1:
				printf("Toggle mode...\n");
				mode = (mode + 1) % 3;
				change_viewmode(mode);
				break;

			case cimg::keyF2:
				printf("Run Watershed...\n");
				app_watershed();
				mode = 1;
				change_viewmode(mode);
				break;

			case cimg::keyF3:
				printf("Choose region...\n");
				app_select();
				mode = 2;
				change_viewmode(mode);
				break;

			case cimg::keyF4: 
				// Find path... 
				// 1. Constrained by region
				// - if no region is selected by the user, it is constrained by
				//   the size of image
				// - input: scaled linechart, region mask
				// - output: tx and ty (translation factor to set a start point)
				printf("Run PSO...\n");
				if (selectedId.size() != 0) {
					(*region).fill(0);
					cimg_forXY(seg_label,x,y) {
						bool bSelected = false;
						for (unsigned int i=0; i<selectedId.size(); i++) {
							if (selectedId[i] == seg_label(x,y)) {
								bSelected = true;
								break;
							}
						}
						if (!bSelected) continue;
						(*region)(x,y) = 1;
					}
				} else {
					(*region).fill(1);
				}

				app_pso(true);
				break;

			case cimg::keyF5:
				printf("Toogle background...\n");
				bgmode = (bgmode + 1) % 2;
				if (bgmode==0)
					img_bg = &img_rgb;
				else if (bgmode==1)
					img_bg = &img_grd;

				change_viewmode(mode);
				break;

			case cimg::keyF6:
				printf("Pick color...\n");
				app_colorpicker();
				refresh_canvas_help();
				break;

			case cimg::keyH:
				printf("Toggle help window...\n");

				onhelp = (onhelp+1) % 2;
				if (onhelp==0) {
					if (!disp_help.is_closed()) {disp_help.close();}
				}
				else {
					if (disp_help.is_closed()) {
						refresh_canvas_help();
						
					}
				}
				break;

			case cimg::keyQ:
				user_scalex += 0.01;
				user_scalex = (user_scalex > 1.0) ? 1.0: user_scalex;
				refresh_canvas_help();
				break;

			case cimg::keyW:
				user_scalex -= 0.01;
				user_scalex = (user_scalex < 0.1) ? 0.1: user_scalex;
				refresh_canvas_help();
				break;

			case cimg::keyA:
				user_scaley += 0.01;
				user_scaley = (user_scaley > 1.0) ? 1.0: user_scaley;
				refresh_canvas_help();
				break;

			case cimg::keyS:
				user_scaley -= 0.01;
				user_scaley = (user_scaley < 0.1) ? 0.1: user_scaley;
				refresh_canvas_help();
				break;

			case cimg::keyE:
				user_tolx += 0.1;
				user_tolx = (user_tolx > 5.0) ? 5.0: user_tolx;
				refresh_canvas_help();
				break;

			case cimg::keyR:
				user_tolx -= 0.1;
				user_tolx = (user_tolx < 0.5) ? 0.5: user_tolx;
				refresh_canvas_help();
				break;

			case cimg::keyD:
				user_toly += 0.1;
				user_toly = (user_toly > 5.0) ? 5.0: user_toly;
				refresh_canvas_help();
				break;

			case cimg::keyF:
				user_toly -= 0.1;
				user_toly = (user_toly < 0.5) ? 0.5: user_toly;
				refresh_canvas_help();
				break;

			case cimg::keyZ:
				edge_threshold += 3.0;
				edge_threshold = (edge_threshold > 255.0) ? 255.0: edge_threshold;
				refresh_canvas_help();
				update_img_grd();
				change_viewmode(mode);
				break;

			case cimg::keyX:
				edge_threshold -= 3.0;
				edge_threshold = (edge_threshold < 0.0) ? 0.0: edge_threshold;
				refresh_canvas_help();
				update_img_grd();
				change_viewmode(mode);
				break;

			case cimg::key1:
				region_factor += 0.1;
				region_factor = (region_factor > 1.0) ? 1.0: region_factor;
				refresh_canvas_help();
				break;

			case cimg::key2:
				region_factor -= 0.1;
				region_factor = (region_factor < 0.0) ? 0.0: region_factor;
				refresh_canvas_help();
				break;

			case cimg::key3:
				edge_factor += 0.1;
				edge_factor = (edge_factor > 1.0) ? 1.0: edge_factor;
				refresh_canvas_help();
				break;

			case cimg::key4:
				edge_factor -= 0.1;
				edge_factor = (edge_factor < 0.0) ? 0.0: edge_factor;
				refresh_canvas_help();
				break;

			case cimg::key5:
				color_factor += 0.1;
				color_factor = (color_factor > 1.0) ? 1.0: color_factor;
				refresh_canvas_help();
				break;

			case cimg::key6:
				color_factor -= 0.1;
				color_factor = (color_factor < 0.0) ? 0.0: color_factor;
				refresh_canvas_help();
				break;

			default:
				break;
			}
			disp_main.set_key();
		}

		if (disp_main.mouse_x()>=0 && disp_main.mouse_y()>=0) {
			;
		}

		canvas_main.display(disp_main);
	}


	delete region;	
	delete linechart;


	/**
	 * PSO
	 * 
	 *	For scale x
	 *		For scale y
	 *			Find tx,ty where maximize object function w.r.t
	 *			the constrained region
	 *		End
	 *  End 
	 */

	return 0;
}